import os
import openpyxl
from paddleocr import PaddleOCR
from concurrent.futures import ThreadPoolExecutor
import queue
import config
from tqdm import tqdm  # 引入进度条库

class SensitiveKeywordsChecker:
    def __init__(self, input_path, output_path):
        self.input_path = input_path
        self.output_path = output_path
        self.result_queue = queue.Queue()  # 线程安全的队列，用于存储结果

        # 准备结果数据表
        self.workbook_ = openpyxl.Workbook()
        self.worksheet_ = self.workbook_.active
        self.worksheet_.append(['序号', '密级', '路径'])
        self.result_seq_ = 0

    def count_jpg_files(self):
        total_jpg_files = 0
        for root, dirs, files in os.walk(self.input_path):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.lower().endswith('.jpg'):
                    total_jpg_files += 1
        return total_jpg_files

    def update_progress(self, current, total):
        # 更新进度条的函数
        #self.progress_bar.update(current)
        #self.progress_bar.set_postfix(total_it=total)
        config.print_console_and_write_log(f'当前进度为：{current},共{total}张图片')


    def process_results(self):
        # 从队列中获取结果并写入Excel
        while not self.result_queue.empty():
            result = self.result_queue.get()
            self.worksheet_.append(result)
        self.workbook_.save(os.path.join(self.output_path, 'result.xlsx'))

    def check_all_jpg_files(self):
        total_jpg_files = self.count_jpg_files()
        logging_info = fr'一共有{total_jpg_files}张jpg图片！'
        config.print_console_and_write_log(logging_info)

        with ThreadPoolExecutor(max_workers=10) as executor:  # 创建线程池，max_workers指定线程数
            futures = []
            self.progress_bar = tqdm(total=total_jpg_files, desc="处理jpg文件")  # 初始化进度条

            for root, dirs, files in os.walk(self.input_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    if file_path.lower().endswith('.jpg'):
                        future = executor.submit(self.check_one_jpg_file, file_path)  # 提交任务到线程池
                        futures.append(future)
                        self.update_progress(len(futures), total_jpg_files)  # 更新进度

            # 等待所有任务完成
            for future in futures:
                try:
                    future.result()  # 获取任务结果，这里实际上不会返回任何值，因为结果已经放入队列
                except Exception as e:
                    config.print_console_and_write_log(f"处理文件时发生错误: {e}")

        # self.progress_bar.close()  # 关闭进度条
        self.process_results()
        config.print_console_and_write_log(f"检查完成，共处理了 {total_jpg_files} 个.jpg文件")

    def check_one_jpg_file(self, image_file_path):
        ocr = PaddleOCR(use_angle_cls=True, lang="ch")
        try:
            result = ocr.ocr(image_file_path, cls=True)
        except Exception as e:
            logging_info = f"OCR 过程中发生错误：{e}"
            self.print_console_and_write_log(logging_info)
            return

        sensitive_words = ['绝密', '机密', '秘密']
        for idx in range(len(result)):
            res = result[idx]
            for line in res:
                target = line[1][0]
                if isinstance(target, str):
                    for word in sensitive_words:
                        if word in target:
                            self.result_seq_ += 1
                            row = [self.result_seq_, image_file_path, word]
                            self.worksheet_.append(row)
                            logging_info = f'找到{word}资料一份，路径为：{image_file_path}'
                            config.print_console_and_write_log(logging_info)
                            return  # 找到敏感词后，可以退出内层循环
